﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Drawing;
using System.Drawing.Imaging;

namespace LowLevelGraphics.Filter
{
    /// <summary>
    /// A mean shift filter is in this case used to segmentate a picture
    /// </summary>
    public class MeanShift : BaseImageFilter
    {
        protected int rad = 7, rad2 = 7;
        protected float radCol, radCol2;
        protected int nPasses, pass;
        protected bool isRGB = true;

        /// <summary>
        /// Initializes a new instance of the <see cref="MeanShift"/> class.
        /// </summary>
        public MeanShift()
        {
        }

        /// <summary>
        /// Initializes a new instance of the <see cref="MeanShift"/> class.
        /// </summary>
        /// <param name="_meanShift">The _mean shift.</param>
        internal MeanShift(MeanShift _meanShift)
        {
        }

        /// <summary>
        /// Execute this Filter
        /// A Debug Statement is written if a concrete Filter is called
        /// </summary>
        /// <param name="_bitmap"></param>
        /// <returns></returns>
        public override UnsafeBitmap Execute(UnsafeBitmap _bitmap)
        {
            if (isRGB)
            {
                FilterColored(_bitmap);
            }
            else
            {
                FilterGrayScale(_bitmap);
            }
            return _bitmap;
        }

        /// <summary>
        /// Filters the colored.
        /// </summary>
        /// <param name="_bitmap">The _bitmap.</param>
        private void FilterColored(UnsafeBitmap _bitmap)
        {
            UnsafeBitmap bitmapCloned = (UnsafeBitmap)_bitmap.Clone();

            int width = _bitmap.Width;
            int height = _bitmap.Height;

            float[, ,] pixelsf = new float[width, height, 3];

            int y = 0;
            int x = 0;
            for (y = 0; y < height; y++)
            {
                byte a = 0;
                byte r = 0;
                byte g = 0;
                byte b = 0;
                bitmapCloned.GetPixelDirect(x, y, out a, out r, out g, out b);

                for (x = 0; x < width; x++)
                {
                    pixelsf[x, y, 0] = 0.299f * r + 0.587f * g + 0.114f * b; // Y
                    pixelsf[x, y, 1] = 0.5957f * r - 0.2744f * g - 0.3212f * b; // I
                    pixelsf[x, y, 2] = 0.2114f * r - 0.5226f * g + 0.3111f * b; // Q
                }
            }

            float shift = 0;
            int iters = 0;
            for (y = 0; y < height; y++)
            {
                for (x = 0; x < width; x++)
                {
                    int xc = x;
                    int yc = y;
                    int xcOld, ycOld;
                    float YcOld, IcOld, QcOld;
                    int pos = y * width + x;
                    float[] yiq = new float[] { pixelsf[x, y, 0], pixelsf[x, y, 1], pixelsf[x, y, 2] };
                    float Yc = yiq[0];
                    float Ic = yiq[1];
                    float Qc = yiq[2];

                    iters = 0;
                    do
                    {
                        xcOld = xc;
                        ycOld = yc;
                        YcOld = Yc;
                        IcOld = Ic;
                        QcOld = Qc;

                        float mx = 0;
                        float my = 0;
                        float mY = 0;
                        float mI = 0;
                        float mQ = 0;
                        int num = 0;

                        for (int ry = -rad; ry <= rad; ry++)
                        {
                            int y2 = yc + ry;
                            if (y2 >= 0 && y2 < height)
                            {
                                for (int rx = -rad; rx <= rad; rx++)
                                {
                                    int x2 = xc + rx;
                                    if (x2 >= 0 && x2 < width)
                                    {
                                        if (ry * ry + rx * rx <= rad2)
                                        {
                                            yiq = new float[] { pixelsf[x2, y2, 0], pixelsf[x2, y2, 1], pixelsf[x2, y2, 2] };

                                            float Y2 = yiq[0];
                                            float I2 = yiq[1];
                                            float Q2 = yiq[2];

                                            float dY = Yc - Y2;
                                            float dI = Ic - I2;
                                            float dQ = Qc - Q2;

                                            if (dY * dY + dI * dI + dQ * dQ <= radCol2)
                                            {
                                                mx += x2;
                                                my += y2;
                                                mY += Y2;
                                                mI += I2;
                                                mQ += Q2;
                                                num++;
                                            }
                                        }
                                    }
                                }
                            }
                        }
                        float num_ = 1f / num;
                        Yc = mY * num_;
                        Ic = mI * num_;
                        Qc = mQ * num_;
                        xc = (int)(mx * num_ + 0.5);
                        yc = (int)(my * num_ + 0.5);
                        int dx = xc - xcOld;
                        int dy = yc - ycOld;
                        float dY2 = Yc - YcOld;
                        float dI2 = Ic - IcOld;
                        float dQ2 = Qc - QcOld;

                        shift = dx * dx + dy * dy + dY2 * dY2 + dI2 * dI2 + dQ2 * dQ2;
                        iters++;
                    }
                    while (shift > 3 && iters < 100);

                    int r_ = (int)(Yc + 0.9563f * Ic + 0.6210f * Qc);
                    int g_ = (int)(Yc - 0.2721f * Ic - 0.6473f * Qc);
                    int b_ = (int)(Yc - 1.1070f * Ic + 1.7046f * Qc);

                    _bitmap.SetPixelDirect(x, y, (byte)r_, (byte)g_, (byte)b_);
                }
            }
            bitmapCloned.Dispose();
        }

        /// <summary>
        /// Filters the gray scale.
        /// </summary>
        /// <param name="_bitmap">The _bitmap.</param>
        protected void FilterGrayScale(UnsafeBitmap _bitmap)
        {
            UnsafeBitmap bitmapCloned = (UnsafeBitmap)_bitmap.Clone();

            int width = _bitmap.Width;
            int height = _bitmap.Height;
            float[,] pixelsf = new float[width, height];

            float shift = 0;
            int iters = 0;

            int y = 0;
            int x = 0;

            for (y = 0; y < height; y++)
            {
                for (x = 0; x < width; x++)
                {
                    int xc = x;
                    int yc = y;
                    int xcOld, ycOld;
                    float YcOld;
                    int pos = y * width + x;
                    float Yc = pixelsf[x, y];

                    iters = 0;
                    do
                    {
                        xcOld = xc;
                        ycOld = yc;
                        YcOld = Yc;

                        float mx = 0;
                        float my = 0;
                        float mY = 0;
                        int num = 0;

                        for (int ry = -rad; ry <= rad; ry++)
                        {
                            int y2 = yc + ry;
                            if (y2 >= 0 && y2 < height)
                            {
                                for (int rx = -rad; rx <= rad; rx++)
                                {
                                    int x2 = xc + rx;
                                    if (x2 >= 0 && x2 < width)
                                    {
                                        if (ry * ry + rx * rx <= rad2)
                                        {

                                            float Y2 = pixelsf[x, y];

                                            float dY = Yc - Y2;

                                            if (dY * dY <= radCol2)
                                            {
                                                mx += x2;
                                                my += y2;
                                                mY += Y2;
                                                num++;
                                            }
                                        }
                                    }
                                }
                            }
                        }
                        float num_ = 1f / num;
                        Yc = mY * num_;
                        xc = (int)(mx * num_ + 0.5);
                        yc = (int)(my * num_ + 0.5);
                        int dx = xc - xcOld;
                        int dy = yc - ycOld;
                        float dY2 = Yc - YcOld;

                        shift = dx * dx + dy * dy + dY2 * dY2;
                        iters++;
                    }
                    while (shift > 3 && iters < 100);

                    _bitmap.SetPixelDirect(x, y, (byte)yc, (byte)yc, (byte)yc);
                }
            }
        }

        /// <summary>
        /// This is used to clone a filter, each filter must implement it's own clone functionality
        /// </summary>
        /// <returns></returns>
        public override object Clone()
        {
            return new MeanShift(this);
        }
    }
}
